import os
import re
from collections import Counter
import csv


def ensure_directory_exists(directory):
    """确保目录存在，如不存在则创建"""
    if not os.path.exists(directory):
        os.makedirs(directory)
        print(f"已创建目录: {directory}")


def count_word_frequency(year):
    """统计指定年份文件中的单词词频"""
    file_path = os.path.join("CET4", f"{year}.txt")

    if not os.path.exists(file_path):
        print(f"错误：找不到文件 {file_path}")
        return None

    # 尝试不同编码读取文件
    for encoding in ['utf-8', 'gbk', 'latin-1']:
        try:
            with open(file_path, 'r', encoding=encoding) as file:
                text = file.read()
            break
        except UnicodeDecodeError:
            continue
    else:
        print(f"错误：无法解码文件 {file_path}")
        return None

    # 提取所有英文单词并转为小写
    words = re.findall(r'\b[a-zA-Z]+\b', text.lower())

    # 统计词频
    word_counts = Counter(words)
    total_words = len(words)

    return word_counts, total_words


def count_multiple_years_frequency(years):
    """统计多个年份文件中的单词词频总和"""
    combined_counts = Counter()
    total_words = 0

    for year in years:
        result = count_word_frequency(year)
        if result:
            year_counts, year_total = result
            combined_counts.update(year_counts)
            total_words += year_total
            print(f"已统计 {year} 年词频，该年份共 {year_total} 个单词")

    if total_words == 0:
        return None

    return combined_counts, total_words


def save_to_csv(word_counts, total_words, output_file):
    """将词频结果保存为CSV文件"""
    # 确保CSV目录存在
    csv_dir = "CSV"
    ensure_directory_exists(csv_dir)

    # 构建完整的输出文件路径
    output_path = os.path.join(csv_dir, output_file)

    # 按词频从高到低排序
    sorted_words = sorted(word_counts.items(), key=lambda x: x[1], reverse=True)

    with open(output_path, 'w', newline='', encoding='utf-8') as csvfile:
        writer = csv.writer(csvfile)
        writer.writerow(["单词", "出现次数", "百分比"])

        for word, count in sorted_words:
            percentage = (count / total_words) * 100
            writer.writerow([word, count, f"{percentage:.4f}%"])

    print(f"词频统计已保存到 {output_path}")
    print(f"共统计了 {len(sorted_words)} 个不同的单词，文本总计 {total_words} 个单词")


def main():
    print("欢迎使用CET4词频统计工具")
    print("请输入要统计的年份，多个年份请用空格分隔（如：22 23 24）")

    while True:
        input_years = input("请输入年份 (输入q退出): ")
        if input_years.lower() == 'q':
            break

        years = input_years.split()
        if not years:
            print("未输入年份，请重新输入")
            continue

        # 生成默认输出文件名
        if len(years) == 1:
            output_file = f"CET4_{years[0]}_词频统计.csv"
        else:
            years_str = "_".join(years)
            output_file = f"CET4_{years_str}_词频统计.csv"

        # 统计词频
        result = count_multiple_years_frequency(years)
        if result:
            word_counts, total_words = result
            save_to_csv(word_counts, total_words, output_file)


if __name__ == "__main__":
    main()